import matplotlib
import numpy as np
import matplotlib.pyplot as plt

def compute_tcb_tangents(P, T=0.0, C=0.0, B=0.0):
    """
    计算每个控制点的切向量（Tangent）.
    参数:
        P : (N, d) 控制点数组，N 为点数，d 为维度（2D/3D）。
        T, C, B : 标量张力、连续性、偏差参数（可为标量或长度为 N 的数组）。
    返回:
        T_out : (N, d) 输出切向量（从点 i 指向 i+1 的方向）。
        T_in  : (N, d) 输入切向量（从点 i 指向 i-1 的方向）。
    """
    N, d = P.shape
    # 统一为数组形式，便于逐点不同参数
    T = np.full(N, T) if np.isscalar(T) else np.asarray(T)
    C = np.full(N, C) if np.isscalar(C) else np.asarray(C)
    B = np.full(N, B) if np.isscalar(B) else np.asarray(B)

    # 前后点差分
    dP_next = np.vstack([P[1:] - P[:-1], np.zeros((1, d))])   # P_{i+1} - P_i
    dP_prev = np.vstack([np.zeros((1, d)), P[1:] - P[:-1]])   # P_i - P_{i-1}

    # 计算系数
    a = (1 - T) * (1 + C) * (1 + B) / 2.0
    b = (1 - T) * (1 - C) * (1 - B) / 2.0
    c = (1 - T) * (1 - C) * (1 + B) / 2.0
    d = (1 - T) * (1 + C) * (1 - B) / 2.0

    # 输出切向量（指向后一个点）
    T_out = a[:, None] * dP_next + b[:, None] * dP_prev
    # 输入切向量（指向前一个点）
    T_in  = c[:, None] * dP_next + d[:, None] * dP_prev

    return T_out, T_in

def hermite_segment(P0, P1, T0, T1, n=100):
    """
    生成两点之间的 Hermite 曲线段。
    参数:
        P0, P1 : 起止点 (d,)
        T0, T1 : 起止切向量 (d,)
        n      : 细分数目
    返回:
        (n, d) 曲线点数组
    """
    t = np.linspace(0, 1, n)[:, None]   # (n,1)
    h00 =  2*t**3 - 3*t**2 + 1
    h10 =      t**3 - 2*t**2 + t
    h01 = -2*t**3 + 3*t**2
    h11 =      t**3 -    t**2
    return h00*P0 + h10*T0 + h01*P1 + h11*T1

def tcb_curve(points, tension=0.0, continuity=0.0, bias=0.0, samples_per_seg=50):
    """
    生成完整的 TCB 曲线。
    参数:
        points          : (N, d) 控制点
        tension, continuity, bias : 标量或长度为 N 的数组
        samples_per_seg  : 每段曲线的采样点数
    返回:
        curve_pts : (M, d) 曲线点集合
    """
    P = np.asarray(points, dtype=float)
    N, d = P.shape
    if N < 2:
        raise ValueError("至少需要两个控制点")

    # 计算切向量
    T_out, T_in = compute_tcb_tangents(P, tension, continuity, bias)

    # 逐段拼接 Hermite 曲线
    curve = []
    for i in range(N - 1):
        seg = hermite_segment(P[i], P[i+1], T_out[i], T_in[i+1], n=samples_per_seg)
        # 去掉每段的最后一个点（除非是最后一段），避免重复
        if i < N - 2:
            seg = seg[:-1]
        curve.append(seg)

    return np.vstack(curve)


def plot_splines():
    matplotlib.rcParams['font.family'] = 'Microsoft YaHei'  # Windows 常用
    # 绘图
    plt.figure(figsize=(8, 5))
    # 曲线
    plt.plot(curve[:, 0], curve[:, 1], 'b-', label='TCB 曲线')
    # 控制点
    plt.plot(ctrl_pts[:, 0], ctrl_pts[:, 1], 'ro', label='控制点')
    # 连线（帮助观察）
    plt.plot(ctrl_pts[:, 0], ctrl_pts[:, 1], 'k--', linewidth=0.5)

    plt.title(f'TCB(T={tension},C={continuity},B={bias})曲线示例')
    plt.legend()
    plt.axis('equal')
    plt.grid(True, linestyle='--', alpha=0.6)
    plt.show()


# ------------------- 示例 -------------------
if __name__ == "__main__":
    # 任意点集（这里演示 2D）
    ctrl_pts = np.array([
        [0, 0],
        [1, 2],
        [3, 3],
        [4, 0],
        [5, -1]
    ])

    # 可自行修改 T、C、B 参数（这里统一为 0，得到 Catmull‑Rom 曲线）
    tension = 0.0
    continuity = 0
    bias = 0.0

    curve = tcb_curve(ctrl_pts, tension, continuity, bias, samples_per_seg=100)
    plot_splines()
